#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt5.QtCore import QObject, QTimer, pyqtSignal

from core.capture import PacketCapturer
from core.parser import PacketParser
from core.analyzer import PacketAnalyzer

class AppController(QObject):
    """
    应用控制器类，负责连接GUI和核心功能模块
    """
    
    # 自定义信号
    packet_captured = pyqtSignal(object)  # 发送新捕获的数据包
    packets_updated = pyqtSignal(list)    # 发送更新后的数据包列表
    statistics_updated = pyqtSignal(dict) # 发送更新后的统计数据
    alert_generated = pyqtSignal(str)     # 发送生成的警报信息
    error_occurred = pyqtSignal(str)      # 发送错误信息
    
    def __init__(self, main_window=None):
        """
        初始化应用控制器
        
        Args:
            main_window: 主窗口对象，用于更新UI
        """
        super().__init__()
        
        self.main_window = main_window
        
        # 创建核心模块实例
        self.capturer = PacketCapturer()
        self.parser = PacketParser()
        self.analyzer = PacketAnalyzer()
        
        # 保存解析后的数据包
        self.parsed_packets = []
        self.display_packets = []  # 用于显示的数据包（可能经过过滤）
        self.display_filter = ""   # 显示过滤器
        
        # 创建定时器，用于定期更新统计信息
        self.stats_timer = QTimer()
        self.stats_timer.timeout.connect(self._update_statistics)
        
        # 连接信号到主窗口
        if main_window:
            self.packets_updated.connect(main_window.update_packet_list)
            self.statistics_updated.connect(main_window.update_statistics)
            self.alert_generated.connect(main_window.display_alert)
            self.error_occurred.connect(main_window.display_error)
            
        # 初始化网络接口列表
        self._init_interfaces()
        
    def _init_interfaces(self):
        """
        初始化网络接口列表
        """
        try:
            interfaces = self.capturer.get_interfaces()
            if self.main_window:
                self.main_window.update_interface_list(interfaces)
        except Exception as e:
            self._on_error(f"初始化网络接口失败: {str(e)}")
            
    def start_capture(self, interface=None, packet_filter=None):
        """
        开始捕获数据包
        
        Args:
            interface: 要捕获的网络接口，可以是接口名称字符串或包含接口信息的字典
            packet_filter: 捕获过滤器字符串
            
        Returns:
            bool: 是否成功开始捕获
        """
        # 重置分析器
        self.analyzer.reset()
        
        # 清空已解析的数据包
        self.parsed_packets = []
        self.display_packets = []
        
        # 更新UI
        self.packets_updated.emit([])
        
        try:
            # 开始捕获
            success = self.capturer.start_capture(
                interface=interface,
                packet_filter=packet_filter,
                callback=self._on_packet_captured,
                error_callback=self._on_error
            )
            
            if success:
                # 启动统计定时器
                self.stats_timer.start(2000)  # 每2秒更新一次统计信息
                self._on_error("捕获已开始", show_popup=False)
            else:
                self._on_error("开始捕获失败")
                
            return success
        except Exception as e:
            self._on_error(f"启动捕获时出错: {str(e)}")
            return False
        
    def stop_capture(self):
        """
        停止捕获数据包
        """
        try:
            self.capturer.stop_capture()
            self.stats_timer.stop()
            
            # 最后更新一次统计信息
            self._update_statistics()
            
            # 确保完整更新UI
            if self.main_window:
                self.packets_updated.emit(self.display_packets)
            
            # 显示总捕获数量，这里设置show_popup=True确保显示弹窗
            packet_count = len(self.parsed_packets)
            total_captured = self.capturer.packet_count if hasattr(self.capturer, 'packet_count') else packet_count
            self._on_error(f"捕获已停止，共捕获 {total_captured} 个数据包，已解析 {packet_count} 个数据包", show_popup=True)
        except Exception as e:
            self._on_error(f"停止捕获时出错: {str(e)}")
        
    def _on_packet_captured(self, packet):
        """
        处理新捕获的数据包
        
        Args:
            packet: 原始数据包对象（scapy packet）
        """
        try:
            # 解析数据包
            parsed_packet = self.parser.parse_packet(packet)
            
            # 添加到列表
            self.parsed_packets.append(parsed_packet)
            
            # 限制内存中保存的解析包数量
            max_parsed_packets = 5000
            if len(self.parsed_packets) > max_parsed_packets:
                # 只保留最新的包
                self.parsed_packets = self.parsed_packets[-max_parsed_packets:]
                # 更新UI
                if len(self.display_packets) > max_parsed_packets:
                    self.display_packets = self.display_packets[-max_parsed_packets:]
                    if self.main_window:
                        self.packets_updated.emit(self.display_packets)
                        self._on_error(f"限制显示最新的 {max_parsed_packets} 个数据包以节省内存", show_popup=False)
            
            # 应用显示过滤器
            if self._apply_display_filter(parsed_packet):
                self.display_packets.append(parsed_packet)
                
                # 发送新数据包到UI（每50个包更新一次UI，避免UI过载）
                if len(self.display_packets) % 50 == 0:
                    if self.main_window:
                        # 使用信号更新UI，避免直接操作
                        self.packets_updated.emit(self.display_packets)
                else:
                    # 单个包更新
                    if self.main_window:
                        self.main_window.packet_list_widget.add_packet(parsed_packet)
                    
            # 分析数据包
            self.analyzer.analyze_packet(parsed_packet)
            
            # 检查是否有新警报
            self._check_for_alerts()
            
        except Exception as e:
            self._on_error(f"处理数据包时出错: {str(e)}")
            
            # 提供更详细的错误信息
            import traceback
            error_details = traceback.format_exc()
            self._on_error(f"详细错误信息: {error_details}", show_popup=False)
            
    def _apply_display_filter(self, packet):
        """
        应用显示过滤器
        
        Args:
            packet: 解析后的数据包
            
        Returns:
            bool: 数据包是否通过过滤器
        """
        if not self.display_filter:
            return True
            
        filter_text = self.display_filter.lower()
        
        # 根据协议过滤
        if "protocol" in packet and filter_text in packet["protocol"].lower():
            return True
            
        # 根据源地址过滤
        if "src" in packet and filter_text in packet["src"].lower():
            return True
            
        # 根据目标地址过滤
        if "dst" in packet and filter_text in packet["dst"].lower():
            return True
            
        # 根据端口过滤
        if "src_port" in packet and filter_text in packet["src_port"]:
            return True
            
        if "dst_port" in packet and filter_text in packet["dst_port"]:
            return True
            
        # 检查协议中是否包含过滤文本
        if "protocols" in packet:
            for protocol in packet["protocols"]:
                if filter_text in protocol.lower():
                    return True
                    
        return False
        
    def apply_display_filter(self, filter_text):
        """
        应用显示过滤器并更新UI
        
        Args:
            filter_text: 过滤器文本
        """
        self.display_filter = filter_text.lower()
        
        # 应用过滤器到所有数据包
        if not filter_text:
            # 如果没有过滤器，显示所有数据包
            self.display_packets = self.parsed_packets.copy()
        else:
            # 应用过滤器
            self.display_packets = [
                packet for packet in self.parsed_packets
                if self._apply_display_filter(packet)
            ]
            
        # 更新UI
        self.packets_updated.emit(self.display_packets)
        
        # 更新过滤后的数据包数量
        if self.main_window:
            self._on_error(f"过滤后显示 {len(self.display_packets)} 个数据包", show_popup=False)
        
    def _update_statistics(self):
        """
        更新统计信息并发送到UI
        """
        try:
            stats_data = self.analyzer.get_summary()
            self.statistics_updated.emit(stats_data)
        except Exception as e:
            self._on_error(f"更新统计信息时出错: {str(e)}", show_popup=False)
        
    def _check_for_alerts(self):
        """
        检查并处理新警报
        """
        try:
            alerts = self.analyzer.get_alerts()
            if alerts:
                # 获取最新的警报
                latest_alert = alerts[-1]
                self.alert_generated.emit(latest_alert)
        except Exception as e:
            self._on_error(f"检查警报时出错: {str(e)}")
            
    def _on_error(self, error_message, show_popup=True):
        """
        处理错误和状态消息
        
        Args:
            error_message: 错误或状态消息
            show_popup: 是否显示弹窗
        """
        try:
            # 发送错误信号到UI
            self.error_occurred.emit(error_message)
            
            # 输出到控制台
            print(f"错误/状态: {error_message}")
        except Exception as e:
            # 出错时至少尝试打印到控制台
            print(f"错误处理失败: {str(e)}")
            print(f"原始消息: {error_message}")
        
    def get_packet_at_index(self, index):
        """
        获取指定索引的数据包
        
        Args:
            index: 数据包索引
            
        Returns:
            dict: 解析后的数据包
        """
        if 0 <= index < len(self.display_packets):
            return self.display_packets[index]
        return None
        
    def load_capture_file(self, filename):
        """
        从文件加载捕获数据
        
        Args:
            filename: pcap文件路径
            
        Returns:
            bool: 是否成功加载
        """
        try:
            # 清空当前数据
            self.parsed_packets = []
            self.display_packets = []
            self.analyzer.reset()
            
            # 更新UI
            self.packets_updated.emit([])
            
            # 定义回调函数，处理从文件加载的每个数据包
            def process_packet(packet):
                parsed_packet = self.parser.parse_packet(packet)
                self.parsed_packets.append(parsed_packet)
                
                if self._apply_display_filter(parsed_packet):
                    self.display_packets.append(parsed_packet)
                    
                self.analyzer.analyze_packet(parsed_packet)
            
            # 从文件加载数据包
            success = self.capturer.load_capture(filename, process_packet)
            
            if success:
                # 更新UI
                self.packets_updated.emit(self.display_packets)
                self._update_statistics()
                
                self._on_error(f"成功从文件加载 {len(self.parsed_packets)} 个数据包")
            else:
                self._on_error(f"加载文件失败: {filename}")
                
            return success
        except Exception as e:
            self._on_error(f"加载捕获文件时出错: {str(e)}")
            return False
        
    def save_capture_file(self, filename):
        """
        保存捕获数据到文件
        
        Args:
            filename: 保存的文件路径
            
        Returns:
            bool: 是否成功保存
        """
        try:
            success = self.capturer.save_capture(filename)
            
            if success:
                self._on_error(f"捕获数据已保存到 {filename}")
            else:
                self._on_error(f"保存捕获数据失败")
                
            return success
        except Exception as e:
            self._on_error(f"保存捕获文件时出错: {str(e)}")
            return False
        
    def export_as_json(self, filename):
        """
        将捕获数据导出为JSON格式
        
        Args:
            filename: 导出的文件路径
            
        Returns:
            bool: 是否成功导出
        """
        try:
            # 使用解析器的导出功能
            success = self.parser.export_to_json(self.parsed_packets, filename)
            
            if success:
                self._on_error(f"数据已导出到 {filename}")
            else:
                self._on_error(f"导出数据失败")
                
            return success
        except Exception as e:
            self._on_error(f"导出JSON时出错: {str(e)}")
            return False
        
    def export_as_csv(self, filename):
        """
        导出捕获数据为CSV格式
        
        Args:
            filename: 导出文件路径
            
        Returns:
            bool: 是否成功导出
        """
        try:
            # 导出为CSV格式，添加UTF-8-BOM标记以确保Windows正确识别编码
            with open(filename, 'w', newline='', encoding='utf-8-sig') as f:
                # 写入CSV头
                f.write("序号,时间,源地址,目标地址,协议,长度,信息\n")
                
                # 写入数据
                for i, packet in enumerate(self.parsed_packets):
                    row = [
                        str(i + 1),
                        packet.get("time", ""),
                        packet.get("src", ""),
                        packet.get("dst", ""),
                        packet.get("protocol", ""),
                        str(packet.get("length", "")),
                        packet.get("info", "")
                    ]
                    # 转义CSV中的特殊字符
                    row = ['"' + str(field).replace('"', '""') + '"' for field in row]
                    f.write(",".join(row) + "\n")
                    
            self._on_error(f"数据已导出到 {filename}")
            return True
        except Exception as e:
            self._on_error(f"导出CSV时出错: {str(e)}")
            return False
        
    def set_port_scan_detection(self, enabled):
        """
        设置是否启用端口扫描检测
        
        Args:
            enabled: 布尔值，True表示启用，False表示禁用
        """
        if self.analyzer:
            self.analyzer.enable_port_scan_detection = enabled
            status = "启用" if enabled else "禁用"
            self._on_error(f"端口扫描检测已{status}", show_popup=False) 